import json
import asyncio

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from starlette.responses import StreamingResponse

from config import HOST, PORT
from mcp_client.mcp_client_sse_demo import chat_with_tools

app = FastAPI()

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class QuestionRequest(BaseModel):
    question: str

@app.post("/LLM")
async def query_handler(request: QuestionRequest):
    question = request.question
    if not question:
        raise HTTPException(status_code=400, detail="缺少问题参数")

    # 流式响应生成器
    async def generate():
        try:
            # 标志位用于区分首次响应
            first_chunk = True

            # 使用请求中的问题调用流式异步生成器
            async for item in chat_with_tools(question):
                if isinstance(item, list):  # 工具调用信息
                    yield f"data: {json.dumps({'type': 'result', 'result': item}, ensure_ascii=False)}\n\n"

                elif isinstance(item, str):  # 流式文本内容
                    if first_chunk:
                        # 发送首次内容标记
                        # yield f"data: {json.dumps({'type': 'content_start'}, ensure_ascii=False)}\n\n"
                        first_chunk = False

                    yield f"data: {json.dumps({'type': 'content', 'content': item}, ensure_ascii=False)}\n\n"

                await asyncio.sleep(0)  # 释放事件循环
                print(item, end="", flush=True)  # 本地终端实时输出)

            # 发送结束标记
            # yield f"data: {json.dumps({'type': 'stream_end'}, ensure_ascii=False)}\n\n"

        except Exception as e:
            error_msg = f"服务异常: {str(e)}"
            yield f"data: {json.dumps({'type': 'error', 'error': error_msg}, ensure_ascii=False)}\n\n"
            print(f"\n[ERROR] {error_msg}")

    # 返回流式响应
    return StreamingResponse(generate(), media_type="text/event-stream")

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host=HOST, port=PORT)